import json
import cv2
import os
import copy
import numpy as np

def add_json(matches, image):
    i = 0
    for match in matches:
        # 在图像上添加文本信息
        text = f"Person: {match['person']['track_id']}: {match['object']['name']}-{match['object']['track_id']}"
        cv2.putText(image, text, (20, 30+i*30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
        i = i + 1
    return image

def load_image(image_folder, filename):
    img_path = os.path.join(image_folder, filename)
    img = cv2.imread(img_path)
    return img

def horizontal_concatenate(img1, img2, direction):
    # 获取两张图片的高度和宽度
    h1, w1, _ = img1.shape
    h2, w2, _ = img2.shape

    # 计算拼接后的图像高度
    max_height = max(h1, h2)
    max_width = max(w1, w2)
    # 创建一个空白图像，用于补全大小不同的地方
    if direction == 0:
        blank1 = np.ones((max_height, w1, 3), dtype=np.uint8) * 255
        blank2 = np.ones((max_height, w2, 3), dtype=np.uint8) * 255
    elif direction == 1:
        blank1 = np.ones((h1, max_width, 3), dtype=np.uint8) * 255
        blank2 = np.ones((h2, max_width, 3), dtype=np.uint8) * 255

    # 将原始图像复制到空白图像上
    blank1[:h1, :w1, :] = img1
    blank2[:h2, :w2, :] = img2

    # 使用hconcat函数水平拼接两个图像
    if direction == 0:
        result = cv2.hconcat([blank1, blank2])
    elif direction == 1:
        result = cv2.vconcat([blank1, blank2])
    return result

def create_collage(person_id, matches, image_folder):
    collage = None
    person_found = False
    object_ids = []
    for match in matches:
        if match['person']['track_id'] == person_id:
            person_found = True
            id = match['object']['track_id']
            if id not in object_ids:
                object_ids.append(id)
    # print(f"{person_id}: {object_ids}")
    if not person_found:
        return None
    person_image = load_image(image_folder, f"{person_id}.jpg")
    linked_image = copy.deepcopy(person_image)
    for object_id in object_ids:
        object_image = load_image(image_folder, f"{object_id}.jpg")
        linked_image = horizontal_concatenate(linked_image, object_image, 0)
    return linked_image
    # return cv2.resize(linked_image, (30, 100))

def generate_frame(result_track, matches_json, images_path):
    person_ids = []
    for track in result_track:
        if track['name'] == 'person':
            person_ids.append(track['track_id'])

    concatenated_image = None
    for person_id in person_ids:
        linked_image = create_collage(person_id, matches_json, images_path)
        if linked_image is None:
            continue
        if concatenated_image is None:
            concatenated_image = linked_image
        else:
            concatenated_image = horizontal_concatenate(concatenated_image, linked_image, 1)
    return concatenated_image
